package xyz.labmem.lipx.server.core.handler

import cn.hutool.core.codec.Base64
import cn.hutool.core.convert.Convert
import cn.hutool.core.date.DateUtil
import cn.hutool.crypto.Mode
import cn.hutool.crypto.Padding
import cn.hutool.crypto.SecureUtil
import cn.hutool.crypto.symmetric.AES
import cn.hutool.crypto.symmetric.SymmetricAlgorithm
import com.alibaba.fastjson2.JSONArray
import io.netty.channel.Channel
import io.netty.channel.ChannelHandlerContext
import io.netty.channel.ChannelInitializer
import io.netty.channel.group.DefaultChannelGroup
import io.netty.channel.socket.SocketChannel
import io.netty.handler.codec.bytes.ByteArrayDecoder
import io.netty.handler.codec.bytes.ByteArrayEncoder
import io.netty.handler.ipfilter.IpFilterRule
import io.netty.handler.ipfilter.IpFilterRuleType
import io.netty.handler.ipfilter.IpSubnetFilterRule
import io.netty.handler.ipfilter.RuleBasedIpFilter
import io.netty.util.concurrent.GlobalEventExecutor
import xyz.labmem.lipx.core.context.ServerContext
import xyz.labmem.lipx.core.context.ServerContext.userList
import xyz.labmem.lipx.core.entity.ClientData
import xyz.labmem.lipx.core.entity.ServerConfig
import xyz.labmem.lipx.core.handler.LabCommonHandler
import xyz.labmem.lipx.core.handler.impl.LocalProxyHandler
import xyz.labmem.lipx.core.handler.impl.RemoteProxyHandler
import xyz.labmem.lipx.core.net.TcpConnection
import xyz.labmem.lipx.core.net.TcpServer
import xyz.labmem.lipx.core.protocol.LabMessage
import xyz.labmem.lipx.core.protocol.LabMessageType
import xyz.labmem.lipx.core.tool.LogInfo
import xyz.labmem.lipx.core.tool.pwdJoinSalt
import xyz.labmem.lipx.labVersion
import java.net.InetSocketAddress
import java.security.SecureRandom
import java.util.concurrent.ConcurrentHashMap


/**
 * @description: do something
 * @author: ltty
 * @date: 2022/5/30 16:38
 */
class LabServerHandler(
    private val server: ServerConfig
) : LabCommonHandler() {

    private var register = false
    private var isReceive = false
    private var port: Int? = null
    private var host: String? = null
    private var ip: String? = null
    private var id: String? = null
    private var aes: AES
    private val aesKey: ByteArray = SecureUtil.generateKey(SymmetricAlgorithm.AES.value, 192).encoded
    private val vi: ByteArray = generateNonce(16)
    private val channelHandlerMap = ConcurrentHashMap<String, LabCommonHandler>()
    private val channelGroup = DefaultChannelGroup(GlobalEventExecutor.INSTANCE)

    private var remoteConnectionServer: TcpServer? = null
    private var receiveConnectionServer: TcpConnection? = null

    init {
        aes = AES(Mode.CTR, Padding.NoPadding, aesKey, vi)
    }

    @Throws(Exception::class)
    override fun channelRead(ctx: ChannelHandlerContext, msg: Any) {
        val message = msg as LabMessage
        ip = (ctx.channel()?.remoteAddress() as InetSocketAddress).address.hostAddress
        if (message.type == LabMessageType.REGISTER) {
            processRegister(message)
        } else if (register) {
            if (message.type === LabMessageType.CONNECTED) { //正向代理连接
                processConnected(message)
            } else if (message.type === LabMessageType.DISCONNECTED) {
                processDisconnected(message)
            } else if (message.type === LabMessageType.DATA) {
                processData(message)
            } else if (message.type === LabMessageType.KEEPALIVE) {
                // 心跳包
                if (isReceive) {
                    ServerContext.userCache.put("receive$ip:$id", userList["receive$ip:$id"])
                } else {
                    ServerContext.userCache.put("$ip:$id", userList["$ip:$id"])
                }
            } else {
                LogInfo.appendLogError("未知请求协议！")
                throw Exception("Unknown type: ${message.type}")
            }
        } else {
            ctx.close()
        }
    }

    @Throws(Exception::class)
    override fun channelInactive(ctx: ChannelHandlerContext) {
        if (register) {
            if (isReceive) {
                channelGroup.close()
                receiveConnectionServer?.channel?.close()
                userList.remove("receive$ip:$id")
            } else {
                remoteConnectionServer?.close()
                userList.remove("$ip:$id")
            }
            LogInfo.appendLog("netty server 已关闭端口：$port")
            ctx.close()
        }
    }

    private fun processConnected(message: LabMessage) {
        receiveConnectionServer = TcpConnection()
        receiveConnectionServer?.connect(host!!, port!!,
            object : ChannelInitializer<SocketChannel>() {
                @Throws(Exception::class)
                override fun initChannel(ch: SocketChannel) {
                    val localProxyHandler =
                        LocalProxyHandler(
                            this@LabServerHandler,
                            message.metaData!!["channelId"].toString(),
                            aes
                        )
                    ch.pipeline().addLast(ByteArrayDecoder(), ByteArrayEncoder(), localProxyHandler)
                    channelHandlerMap[message.metaData!!["channelId"].toString()] = localProxyHandler
                    channelGroup.add(ch)
                }
            })
    }

    /**
     * REGISTER
     */
    private fun processRegister(message: LabMessage) {
        val metaData = HashMap<String, Any>()
        val password: String = message.metaData!!["password"].toString()
        val cwls = message.metaData!!["wls"] as JSONArray
        if (server.wls.isEmpty() || server.wls.contains(ip)) {
            if (pwdJoinSalt(server.password) != password) {
                metaData["success"] = false
                metaData["reason"] = "Token is wrong"
            } else {
                try {
                    isReceive = message.metaData!!["isReceive"] as Boolean
                    metaData["isReceive"] = isReceive
                    id = message.metaData!!["id"] as String
                    if (isReceive) {
                        host = message.metaData!!["proxyHost"].toString()
                        port = Convert.toInt(message.metaData!!["proxyPort"])
                        //创建密钥
                        metaData["success"] = true
                        register = true
                        userList["receive$ip:$id"] = ClientData().apply c@{
                            name = message.metaData!!["name"].toString()
                            ip = this@LabServerHandler.ip
                            id = this@LabServerHandler.id
                            host = this@LabServerHandler.host
                            this.port = this@LabServerHandler.port
                            connectTime = DateUtil.now()
                            connectChannel = this@LabServerHandler
                            secretKey = aesKey
                            isReceive = true
                            metaData["SecretKey"] = Base64.encode(aesKey)
                            metaData["vi"] = Base64.encode(vi)
                        }
                        LogInfo.appendLog("客户端[$ip]连接成功，监听正向代理端口： $host:$port")
                    } else {
                        var success = false
                        val port = message.metaData!!["port"] as Int
                        val thisHandler = this
                        remoteConnectionServer = TcpServer()
                        remoteConnectionServer!!.bind(port, object : ChannelInitializer<SocketChannel>() {
                            @Throws(Exception::class)
                            override fun initChannel(ch: SocketChannel) {
                                ch.pipeline().apply {
                                    if (cwls.isNotEmpty()) {
                                        val rejectAll = object : IpFilterRule {
                                            override fun matches(remoteAddress: InetSocketAddress): Boolean = true
                                            override fun ruleType(): IpFilterRuleType = IpFilterRuleType.REJECT
                                        }
                                        addLast(
                                            "ipFilter",
                                            RuleBasedIpFilter(* cwls.map {
                                                IpSubnetFilterRule(
                                                    it.toString(),
                                                    32,
                                                    IpFilterRuleType.ACCEPT
                                                )
                                            }.toTypedArray(), rejectAll)
                                        )
                                    }
                                    addLast(
                                        ByteArrayDecoder(),
                                        ByteArrayEncoder(),
                                        RemoteProxyHandler(thisHandler, port, aes)
                                    )
                                }
                                ServerContext.channels.add(ch)
                            }
                        }) {
                            this.port = port
                            success = true
                        }
                        if (success) {
                            //创建密钥
                            metaData["success"] = true
                            register = true
                            userList["$ip:$id"] = ClientData().apply c@{
                                name = message.metaData!!["name"].toString()
                                ip = this@LabServerHandler.ip
                                id = this@LabServerHandler.id
                                this.port = this@LabServerHandler.port
                                connectTime = DateUtil.now()
                                connectChannel = this@LabServerHandler
                                secretKey = aesKey
                                metaData["SecretKey"] = Base64.encode(aesKey)
                                metaData["vi"] = Base64.encode(vi)
                            }
                            LogInfo.appendLog("客户端[$ip]注册成功，启动代理端口： $port")
                        }
                    }
                    metaData["version"] = labVersion
                } catch (e: Exception) {
                    metaData["success"] = false
                    metaData["reason"] = e.message!!
                    e.printStackTrace()
                }
            }
        } else {
            metaData["success"] = false
            metaData["reason"] = "拒绝连接！"
            LogInfo.appendLogError("客户端[$ip]不在白名单，已拒绝连接！")
        }
        LabMessage().apply {
            type = LabMessageType.REGISTER_RESULT
            this.metaData = metaData
            ctx?.writeAndFlush(this)
        }
        if (!register) {
            LogInfo.appendLogError("客户端[$ip]注册失败:${metaData["reason"]}")
            ctx?.close()
        }
    }

    /**
     * DATA
     */
    private fun processData(message: LabMessage) {
        if (isReceive) {
            val channelId: String = message.metaData!!["channelId"].toString()
            val handler = channelHandlerMap[channelId]
            if (handler != null) {
                val ctx = handler.ctx
                ctx?.writeAndFlush(aes.decrypt(message.data) ?: message.data)
            }
        } else {
            ServerContext.channels.writeAndFlush(
                aes.decrypt(message.data)
            ) { channel: Channel ->
                channel.id().asLongText() == message.metaData!!["channelId"]
            }
        }
    }

    /**
     * DISCONNECTED
     */
    private fun processDisconnected(message: LabMessage) {
        if (isReceive) {
            val channelId: String = message.metaData!!["channelId"].toString()
            val handler = channelHandlerMap[channelId]
            if (handler != null) {
                handler.ctx?.close()
                channelHandlerMap.remove(channelId)
            }
        } else {
            ServerContext.channels.close { channel: Channel ->
                channel.id().asLongText() == message.metaData!!["channelId"]
            }
        }
    }

    private fun generateNonce(size: Int): ByteArray {
        val nonce = ByteArray(size)
        val random = SecureRandom()
        random.nextBytes(nonce)
        return nonce
    }

}